1378 lines
51 KiB
Rust
1378 lines
51 KiB
Rust
use std::ffi::OsStr;
|
|
use std::fs::File;
|
|
use std::io::Read;
|
|
use std::ops::Deref;
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::{Arc, Mutex as StdMutex};
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
use tokio::sync::Mutex;
|
|
|
|
use crate::db::agent_driver::{
|
|
validate_dameng_java_system_properties, AgentDriverClient, AgentLaunchSpec, AgentMethod, AgentRuntimeClient,
|
|
};
|
|
use crate::models::connection::DatabaseType;
|
|
|
|
pub const DEFAULT_JRE_KEY: &str = "21";
|
|
pub const SQLITE_WORKER_DRIVER_KEY: &str = "sqlite-worker";
|
|
pub const SQLITE_WORKER_NATIVE_PLATFORMS: &[&str] = &["linux-x64", "linux-aarch64"];
|
|
pub const DOWNLOAD_CACHE_DIR_NAME: &str = "download-cache";
|
|
pub const DOWNLOAD_CACHE_MAX_AGE_DAYS: u64 = 7;
|
|
|
|
pub type OperationLockTable = StdMutex<std::collections::HashMap<String, Arc<Mutex<()>>>>;
|
|
|
|
pub struct OperationLockHandle<'a> {
|
|
table: &'a OperationLockTable,
|
|
key: String,
|
|
lock: Arc<Mutex<()>>,
|
|
}
|
|
|
|
impl<'a> OperationLockHandle<'a> {
|
|
pub fn new(table: &'a OperationLockTable, key: &str, lock: Arc<Mutex<()>>) -> Self {
|
|
Self { table, key: key.to_string(), lock }
|
|
}
|
|
}
|
|
|
|
impl Deref for OperationLockHandle<'_> {
|
|
type Target = Mutex<()>;
|
|
|
|
fn deref(&self) -> &Self::Target {
|
|
&self.lock
|
|
}
|
|
}
|
|
|
|
impl Drop for OperationLockHandle<'_> {
|
|
fn drop(&mut self) {
|
|
let Ok(mut locks) = self.table.lock() else {
|
|
return;
|
|
};
|
|
let should_remove =
|
|
locks.get(&self.key).is_some_and(|lock| Arc::ptr_eq(lock, &self.lock) && Arc::strong_count(lock) == 2);
|
|
if should_remove {
|
|
locks.remove(&self.key);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Cooperative cancellation token for an agent driver install/upgrade.
|
|
///
|
|
/// Mirrors the updater's `DownloadCancellation` (src-tauri/src/commands/update.rs)
|
|
/// so the same `watch<bool>` pattern can abort an in-flight JRE/driver download.
|
|
#[derive(Debug)]
|
|
pub struct AgentInstallCancellation {
|
|
canceled: tokio::sync::watch::Sender<bool>,
|
|
}
|
|
|
|
impl AgentInstallCancellation {
|
|
fn new() -> Self {
|
|
let (canceled, _) = tokio::sync::watch::channel(false);
|
|
Self { canceled }
|
|
}
|
|
|
|
pub fn cancel(&self) {
|
|
let _ = self.canceled.send_replace(true);
|
|
}
|
|
|
|
pub fn is_cancelled(&self) -> bool {
|
|
*self.canceled.borrow()
|
|
}
|
|
|
|
/// Completes once cancellation is requested (immediately if already cancelled).
|
|
pub async fn cancelled(&self) {
|
|
let mut receiver = self.canceled.subscribe();
|
|
if *receiver.borrow_and_update() {
|
|
return;
|
|
}
|
|
while receiver.changed().await.is_ok() {
|
|
if *receiver.borrow_and_update() {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn default_jre_key() -> String {
|
|
DEFAULT_JRE_KEY.to_string()
|
|
}
|
|
|
|
fn strip_utf8_bom(value: &str) -> &str {
|
|
value.strip_prefix('\u{feff}').unwrap_or(value)
|
|
}
|
|
|
|
fn is_valid_jar_file(path: &Path) -> bool {
|
|
if !path.is_file() {
|
|
return false;
|
|
}
|
|
// Unzipping the manifest goes through zlib's inflate, which reserves a large
|
|
// buffer on the stack. Run it on a dedicated thread so it gets a full, fresh
|
|
// stack instead of whatever headroom is left at the bottom of a deep async
|
|
// call chain (that chain can already consume most of a tokio worker stack).
|
|
std::thread::scope(|scope| scope.spawn(|| read_jar_manifest(path)).join().unwrap_or(false))
|
|
}
|
|
|
|
fn read_jar_manifest(path: &Path) -> bool {
|
|
let Some(file) = File::open(path).ok() else {
|
|
return false;
|
|
};
|
|
let Some(mut archive) = zip::ZipArchive::new(file).ok() else {
|
|
return false;
|
|
};
|
|
let Some(mut manifest) = archive.by_name("META-INF/MANIFEST.MF").ok() else {
|
|
return false;
|
|
};
|
|
let mut manifest_text = String::new();
|
|
manifest.read_to_string(&mut manifest_text).is_ok() && manifest_text.contains("Main-Class:")
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AgentRegistry {
|
|
#[serde(default)]
|
|
pub jre: Option<JreInfo>,
|
|
#[serde(default)]
|
|
pub jres: std::collections::HashMap<String, JreInfo>,
|
|
pub drivers: std::collections::HashMap<String, DriverInfo>,
|
|
}
|
|
|
|
impl AgentRegistry {
|
|
pub fn resolve_jre(&self, key: &str) -> Option<&JreInfo> {
|
|
if !self.jres.is_empty() {
|
|
return self.jres.get(key);
|
|
}
|
|
if key == DEFAULT_JRE_KEY {
|
|
self.jre.as_ref()
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::fs;
|
|
|
|
fn test_manager(name: &str) -> AgentManager {
|
|
let dir = std::env::temp_dir().join(format!("dbx-agent-manager-{name}-{}", uuid::Uuid::new_v4()));
|
|
AgentManager::new_with_base_dir(dir)
|
|
}
|
|
|
|
fn touch(path: &Path) {
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).unwrap();
|
|
}
|
|
fs::write(path, b"test executable").unwrap();
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
|
|
let mut permissions = fs::metadata(path).unwrap().permissions();
|
|
permissions.set_mode(0o755);
|
|
fs::set_permissions(path, permissions).unwrap();
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cleanup_pending_jre_removes_stash_dirs_and_persists() {
|
|
let manager = test_manager("pending-cleanup");
|
|
std::fs::create_dir_all(manager.base_dir()).unwrap();
|
|
let stash = manager.base_dir().join("jre-21.old-1700000000-deadbeef");
|
|
std::fs::create_dir_all(&stash).unwrap();
|
|
std::fs::write(stash.join("dummy"), b"x").unwrap();
|
|
|
|
let mut state = AgentState::default();
|
|
state.pending_jre_cleanup.push(stash.clone());
|
|
manager.save_state(&state).unwrap();
|
|
|
|
// Re-create the manager (simulates app restart).
|
|
let manager2 = AgentManager::new_with_base_dir(manager.base_dir().clone());
|
|
|
|
assert!(!stash.exists(), "stash dir should be removed");
|
|
let after = manager2.load_state();
|
|
assert!(after.pending_jre_cleanup.is_empty(), "cleanup should drain the list on success");
|
|
}
|
|
|
|
#[test]
|
|
fn cleanup_orphan_jre_dirs_removes_unrecorded_stash() {
|
|
let manager = test_manager("orphan-cleanup");
|
|
std::fs::create_dir_all(manager.base_dir()).unwrap();
|
|
let orphan = manager.base_dir().join("jre-21.old-1234567890-cafe");
|
|
std::fs::create_dir_all(&orphan).unwrap();
|
|
|
|
// Re-create manager — it should sweep orphans even without state.
|
|
let _manager2 = AgentManager::new_with_base_dir(manager.base_dir().clone());
|
|
|
|
assert!(!orphan.exists(), "orphan stash should be swept");
|
|
}
|
|
|
|
#[test]
|
|
fn cleanup_skips_active_jre_dir() {
|
|
let manager = test_manager("orphan-skip-active");
|
|
std::fs::create_dir_all(manager.base_dir()).unwrap();
|
|
let active = manager.jre_dir(DEFAULT_JRE_KEY); // jre-21
|
|
std::fs::create_dir_all(&active).unwrap();
|
|
|
|
let _manager2 = AgentManager::new_with_base_dir(manager.base_dir().clone());
|
|
|
|
assert!(active.exists(), "active jre-<key> dir must not be touched (no .old- in name)");
|
|
}
|
|
|
|
#[test]
|
|
fn agent_state_back_compat_without_pending_jre_cleanup() {
|
|
// Old state JSON without pending_jre_cleanup must still deserialize.
|
|
let json = r#"{
|
|
"jre_versions": {},
|
|
"installed_drivers": {},
|
|
"java_runtime": {"mode": "managed"}
|
|
}"#;
|
|
let state: AgentState = serde_json::from_str(json).expect("deserialize legacy state");
|
|
assert!(state.pending_jre_cleanup.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn loads_agent_state_with_utf8_bom() {
|
|
let manager = test_manager("state-utf8-bom");
|
|
fs::create_dir_all(manager.base_dir()).unwrap();
|
|
fs::write(
|
|
manager.state_path(),
|
|
b"\xEF\xBB\xBF{\"installed_drivers\":{\"kafka\":{\"version\":\"0.1.4\",\"installed_at\":\"now\",\"jre\":\"21\"}}}",
|
|
)
|
|
.unwrap();
|
|
|
|
let state = manager.load_state();
|
|
|
|
assert_eq!(state.installed_drivers.get("kafka").map(|driver| driver.version.as_str()), Some("0.1.4"));
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_managed_java_runtime_by_default() {
|
|
let manager = test_manager("managed");
|
|
let java = manager.jre_java_path(DEFAULT_JRE_KEY);
|
|
touch(&java);
|
|
|
|
let state = AgentState::default();
|
|
|
|
assert_eq!(manager.resolve_java_runtime(&state, DEFAULT_JRE_KEY).unwrap(), java);
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_custom_java_runtime_when_configured() {
|
|
let manager = test_manager("custom");
|
|
let custom_java = manager.base_dir().join("custom").join("bin").join("java");
|
|
touch(&custom_java);
|
|
let state = AgentState {
|
|
java_runtime: JavaRuntimeConfig {
|
|
mode: JavaRuntimeMode::Custom,
|
|
custom_java_path: Some(custom_java.to_string_lossy().to_string()),
|
|
},
|
|
..AgentState::default()
|
|
};
|
|
|
|
assert_eq!(manager.resolve_java_runtime(&state, DEFAULT_JRE_KEY).unwrap(), custom_java);
|
|
}
|
|
|
|
#[test]
|
|
fn rejects_missing_custom_java_runtime() {
|
|
let manager = test_manager("missing-custom");
|
|
let state = AgentState {
|
|
java_runtime: JavaRuntimeConfig {
|
|
mode: JavaRuntimeMode::Custom,
|
|
custom_java_path: Some(manager.base_dir().join("missing-java").to_string_lossy().to_string()),
|
|
},
|
|
..AgentState::default()
|
|
};
|
|
|
|
let err = manager.resolve_java_runtime(&state, DEFAULT_JRE_KEY).unwrap_err();
|
|
|
|
assert!(err.contains("Custom Java runtime does not exist"));
|
|
}
|
|
|
|
#[test]
|
|
fn stores_configured_app_version_for_agent_handshake() {
|
|
let dir = std::env::temp_dir().join(format!("dbx-agent-manager-version-{}", uuid::Uuid::new_v4()));
|
|
let manager = AgentManager::new_with_base_dir_and_app_version(dir, "0.5.13");
|
|
|
|
assert_eq!(manager.agent_app_version(), "0.5.13");
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_system_java_runtime_from_path() {
|
|
let manager = test_manager("system");
|
|
let system_java = manager.base_dir().join("bin").join(java_executable_name());
|
|
touch(&system_java);
|
|
let path = std::env::join_paths([system_java.parent().unwrap()]).unwrap();
|
|
|
|
assert_eq!(resolve_system_java_path(None, Some(path.as_os_str())).unwrap(), system_java);
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_system_java_runtime_from_java_home() {
|
|
let manager = test_manager("system-home");
|
|
let java_home = manager.base_dir().join("jdk");
|
|
let system_java = java_home.join("bin").join(java_executable_name());
|
|
touch(&system_java);
|
|
// PATH deliberately points somewhere else: JAVA_HOME is the explicit
|
|
// declaration and must win.
|
|
let other_bin = manager.base_dir().join("other-bin");
|
|
let other_java = other_bin.join(java_executable_name());
|
|
touch(&other_java);
|
|
let path = std::env::join_paths([&other_bin]).unwrap();
|
|
|
|
assert_eq!(resolve_system_java_path(Some(java_home.as_os_str()), Some(path.as_os_str())).unwrap(), system_java);
|
|
}
|
|
|
|
#[test]
|
|
fn stale_java_home_falls_back_to_path() {
|
|
let manager = test_manager("system-stale-home");
|
|
let system_java = manager.base_dir().join("bin").join(java_executable_name());
|
|
touch(&system_java);
|
|
let path = std::env::join_paths([system_java.parent().unwrap()]).unwrap();
|
|
let dangling = manager.base_dir().join("removed-jdk");
|
|
|
|
assert_eq!(resolve_system_java_path(Some(dangling.as_os_str()), Some(path.as_os_str())).unwrap(), system_java);
|
|
assert!(resolve_system_java_path(Some(dangling.as_os_str()), None).is_none());
|
|
assert!(resolve_system_java_path(Some(OsStr::new("")), None).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn system_java_missing_error_names_the_checked_sources() {
|
|
let dangling = std::env::temp_dir().join("dbx-missing-jdk");
|
|
let path = std::env::join_paths([std::env::temp_dir()]).unwrap();
|
|
|
|
let message = system_java_missing_error(Some(dangling.as_os_str()), Some(path.as_os_str()));
|
|
assert!(message.contains(&dangling.display().to_string()), "{message}");
|
|
assert!(message.contains(java_executable_name()), "{message}");
|
|
assert!(message.contains("no java in the 1 PATH entries checked"), "{message}");
|
|
assert!(message.contains("restart dbx"), "{message}");
|
|
|
|
let unset = system_java_missing_error(None, None);
|
|
assert!(unset.contains("JAVA_HOME is not set"), "{unset}");
|
|
assert!(unset.contains("PATH is not set"), "{unset}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn runtime_gateway_returns_existing_missing_driver_error() {
|
|
let manager = test_manager("missing-driver");
|
|
|
|
let err = match manager.spawn(&DatabaseType::H2, None).await {
|
|
Ok(_) => panic!("missing driver should fail"),
|
|
Err(err) => err,
|
|
};
|
|
|
|
assert_eq!(err, "h2 driver is not installed. Please install it from the Driver Manager.");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn runtime_gateway_returns_existing_missing_java_error() {
|
|
let manager = test_manager("missing-java");
|
|
let jar = manager.driver_jar_path("h2");
|
|
touch(&jar);
|
|
|
|
let err = match manager.spawn(&DatabaseType::H2, None).await {
|
|
Ok(_) => panic!("missing Java runtime should fail"),
|
|
Err(err) => err,
|
|
};
|
|
|
|
assert_eq!(err, "JRE 21 runtime is not installed. Please install it from the Driver Manager.");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn runtime_gateway_resolves_profile_specific_keys() {
|
|
let manager = test_manager("profile-key");
|
|
|
|
assert_eq!(AgentManager::db_type_to_agent_key(&DatabaseType::Oracle, Some("oracle-10g")), Some("oracle"));
|
|
assert_eq!(AgentManager::db_type_to_agent_key(&DatabaseType::Oracle, Some("oracle-legacy")), Some("oracle"));
|
|
assert_eq!(AgentManager::db_type_to_agent_key(&DatabaseType::Oracle, None), Some("oracle"));
|
|
assert_eq!(AgentManager::db_type_to_agent_key(&DatabaseType::Gbase, Some("gbase8s")), Some("gbase8s"));
|
|
assert_eq!(AgentManager::db_type_to_agent_key(&DatabaseType::Gbase, None), Some("gbase8a"));
|
|
manager.stop_daemon_by_key("oracle").await;
|
|
manager.stop_daemon_by_key("gbase8s").await;
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_native_agent_launch_when_agent_executable_exists() {
|
|
let manager = test_manager("native-agent");
|
|
let native = manager.driver_native_path("dameng");
|
|
touch(&native);
|
|
|
|
let launch = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "dameng", DEFAULT_JRE_KEY)
|
|
.expect("native launch should resolve");
|
|
|
|
assert_eq!(launch.program, native);
|
|
assert_eq!(launch.args, Vec::<String>::new());
|
|
assert_eq!(launch.working_dir.as_deref(), Some(manager.driver_dir("dameng").as_path()));
|
|
}
|
|
|
|
#[test]
|
|
fn hive_native_launch_ignores_stale_legacy_jar_without_jre() {
|
|
let manager = test_manager("hive-native-over-legacy-jar");
|
|
let native = manager.driver_native_path("hive");
|
|
touch(&native);
|
|
touch(&manager.driver_jar_path("hive"));
|
|
|
|
let launch = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "hive", DEFAULT_JRE_KEY)
|
|
.expect("Hive native launch should not require a JRE");
|
|
|
|
assert_eq!(launch.program, native);
|
|
assert_eq!(launch.args, Vec::<String>::new());
|
|
assert_eq!(launch.working_dir.as_deref(), Some(manager.driver_dir("hive").as_path()));
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_relative_native_agent_launch_with_absolute_paths() {
|
|
let base_dir =
|
|
PathBuf::from("target").join(format!("dbx-agent-manager-relative-native-{}", uuid::Uuid::new_v4()));
|
|
let manager = AgentManager::new_with_base_dir(base_dir.clone());
|
|
let native = manager.driver_native_path("oracle");
|
|
touch(&native);
|
|
let expected_program = native.canonicalize().unwrap();
|
|
let expected_working_dir = manager.driver_dir("oracle").canonicalize().unwrap();
|
|
|
|
let launch = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "oracle", DEFAULT_JRE_KEY)
|
|
.expect("relative native launch should resolve");
|
|
|
|
assert!(launch.program.is_absolute());
|
|
assert_eq!(launch.program, expected_program);
|
|
assert_eq!(launch.args, Vec::<String>::new());
|
|
assert_eq!(launch.working_dir.as_deref(), Some(expected_working_dir.as_path()));
|
|
|
|
fs::remove_dir_all(base_dir).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn relative_native_agent_launch_keeps_missing_driver_error() {
|
|
let base_dir =
|
|
PathBuf::from("target").join(format!("dbx-agent-manager-relative-missing-{}", uuid::Uuid::new_v4()));
|
|
let manager = AgentManager::new_with_base_dir(base_dir);
|
|
|
|
let error = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "oracle", DEFAULT_JRE_KEY)
|
|
.expect_err("missing native agent should fail");
|
|
|
|
assert_eq!(error, "oracle driver is not installed. Please install it from the Driver Manager.");
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_manifest_agent_launch_with_driver_dir_templates() {
|
|
let manager = test_manager("manifest-agent");
|
|
let driver_dir = manager.driver_dir("dameng-go");
|
|
fs::create_dir_all(driver_dir.join("bin")).unwrap();
|
|
fs::write(
|
|
manager.driver_launch_config_path("dameng-go"),
|
|
r#"{
|
|
"command": "bin/dameng-agent",
|
|
"args": ["--config", "{driver_dir}/config.json"],
|
|
"working_dir": "{driver_dir}"
|
|
}"#,
|
|
)
|
|
.unwrap();
|
|
|
|
let launch = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "dameng-go", DEFAULT_JRE_KEY)
|
|
.expect("manifest launch should resolve");
|
|
|
|
assert_eq!(launch.program, driver_dir.join("bin").join("dameng-agent"));
|
|
assert_eq!(
|
|
launch.args,
|
|
vec!["--config".to_string(), driver_dir.join("config.json").to_string_lossy().to_string()]
|
|
);
|
|
assert_eq!(launch.working_dir.as_deref(), Some(driver_dir.as_path()));
|
|
}
|
|
|
|
#[test]
|
|
fn resolves_manifest_agent_launch_with_utf8_bom() {
|
|
let manager = test_manager("manifest-agent-utf8-bom");
|
|
let driver_dir = manager.driver_dir("kafka");
|
|
fs::create_dir_all(&driver_dir).unwrap();
|
|
fs::write(
|
|
manager.driver_launch_config_path("kafka"),
|
|
b"\xEF\xBB\xBF{\"command\":\"java\",\"args\":[\"-jar\",\"{driver_dir}/agent.jar\"]}",
|
|
)
|
|
.unwrap();
|
|
|
|
let launch = manager
|
|
.resolve_agent_launch_spec(&AgentState::default(), "kafka", DEFAULT_JRE_KEY)
|
|
.expect("manifest launch with UTF-8 BOM should resolve");
|
|
|
|
assert_eq!(launch.program, PathBuf::from("java"));
|
|
assert_eq!(launch.args, vec!["-jar".to_string(), format!("{}/agent.jar", driver_dir.to_string_lossy())]);
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct JreInfo {
|
|
pub version: String,
|
|
pub platforms: std::collections::HashMap<String, ArtifactInfo>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DriverInfo {
|
|
pub version: String,
|
|
pub label: String,
|
|
pub min_app_version: String,
|
|
#[serde(default)]
|
|
pub jar: Option<ArtifactInfo>,
|
|
#[serde(default)]
|
|
pub native: std::collections::HashMap<String, ArtifactInfo>,
|
|
#[serde(default = "default_jre_key")]
|
|
pub jre: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ArtifactInfo {
|
|
pub url: String,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub sha256: Option<String>,
|
|
pub size: u64,
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub format: Option<ArtifactFormat>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum ArtifactFormat {
|
|
TarZstd,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
|
pub struct AgentState {
|
|
#[serde(default)]
|
|
pub jre_version: Option<String>,
|
|
#[serde(default)]
|
|
pub jre_versions: std::collections::HashMap<String, String>,
|
|
#[serde(default)]
|
|
pub installed_drivers: std::collections::HashMap<String, InstalledDriver>,
|
|
#[serde(default)]
|
|
pub java_runtime: JavaRuntimeConfig,
|
|
/// Old JRE directories that could not be deleted in-place during a
|
|
/// reinstall on Windows; renamed aside (`<name>.old-<ts>-<rand>`) and
|
|
/// cleaned up best-effort on next `AgentManager::new`.
|
|
#[serde(default)]
|
|
pub pending_jre_cleanup: Vec<PathBuf>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct InstalledDriver {
|
|
pub version: String,
|
|
pub installed_at: String,
|
|
#[serde(default = "default_jre_key")]
|
|
pub jre: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AgentLaunchConfig {
|
|
pub command: String,
|
|
#[serde(default)]
|
|
pub args: Vec<String>,
|
|
#[serde(default)]
|
|
pub working_dir: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum JavaRuntimeMode {
|
|
#[default]
|
|
Managed,
|
|
System,
|
|
Custom,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
|
pub struct JavaRuntimeConfig {
|
|
#[serde(default)]
|
|
pub mode: JavaRuntimeMode,
|
|
#[serde(default)]
|
|
pub custom_java_path: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AgentDriverInfo {
|
|
pub db_type: String,
|
|
pub label: String,
|
|
pub version: String,
|
|
pub size: u64,
|
|
pub installed: bool,
|
|
pub installed_version: Option<String>,
|
|
pub update_available: bool,
|
|
pub requires_java_runtime: bool,
|
|
pub jre: String,
|
|
pub jre_installed: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DriverStoreUsageItem {
|
|
pub id: String,
|
|
pub bytes: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DriverStoreUsage {
|
|
pub total_bytes: u64,
|
|
pub jre_bytes: u64,
|
|
pub agent_driver_bytes: u64,
|
|
#[serde(default)]
|
|
pub download_cache_bytes: u64,
|
|
pub jdbc_plugin_bytes: u64,
|
|
pub jdbc_driver_bytes: u64,
|
|
pub jres: Vec<DriverStoreUsageItem>,
|
|
pub agent_drivers: Vec<DriverStoreUsageItem>,
|
|
}
|
|
|
|
pub struct AgentManager {
|
|
base_dir: PathBuf,
|
|
app_version: String,
|
|
pub daemons: Mutex<std::collections::HashMap<String, AgentDriverClient>>,
|
|
pub connection_runtimes: Mutex<
|
|
std::collections::HashMap<String, std::sync::Arc<tokio::sync::OnceCell<std::sync::Arc<AgentRuntimeClient>>>>,
|
|
>,
|
|
/// Serializes `load_state` → modify → `save_state` to prevent lost updates
|
|
/// when multiple driver installs run concurrently.
|
|
pub state_lock: StdMutex<()>,
|
|
/// Per-JRE-key install locks so that concurrent driver installs sharing the
|
|
/// same JRE download it only once (DCL pattern: lock → re-check installed → download).
|
|
pub jre_install_locks: OperationLockTable,
|
|
/// Per-driver locks serialize install, import, and uninstall operations
|
|
/// targeting the same on-disk agent files.
|
|
pub driver_operation_locks: OperationLockTable,
|
|
/// Driver operations may run concurrently, but JRE replacement/removal
|
|
/// must exclude them until their dependent driver state is persisted.
|
|
pub installation_operation_lock: tokio::sync::RwLock<()>,
|
|
/// Per-operation cancellation tokens keyed by operation id (single installs
|
|
/// use `install:<op>`, batches `batch:<op>` and `batch:<op>:<db>`). Downloads
|
|
/// observe the exact token threaded to them to abort promptly.
|
|
pub install_cancellations: tokio::sync::Mutex<std::collections::HashMap<String, Arc<AgentInstallCancellation>>>,
|
|
}
|
|
|
|
impl Default for AgentManager {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl AgentManager {
|
|
pub fn new() -> Self {
|
|
let home =
|
|
std::env::var(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).unwrap_or_else(|_| ".".to_string());
|
|
Self::new_with_base_dir(PathBuf::from(home).join(".dbx").join("agents"))
|
|
}
|
|
|
|
pub fn new_with_base_dir(base_dir: PathBuf) -> Self {
|
|
Self::new_with_base_dir_and_app_version(base_dir, env!("CARGO_PKG_VERSION"))
|
|
}
|
|
|
|
pub fn new_with_base_dir_and_app_version(base_dir: PathBuf, app_version: impl Into<String>) -> Self {
|
|
let mgr = Self {
|
|
base_dir,
|
|
app_version: app_version.into(),
|
|
daemons: Mutex::new(std::collections::HashMap::new()),
|
|
connection_runtimes: Mutex::new(std::collections::HashMap::new()),
|
|
state_lock: StdMutex::new(()),
|
|
jre_install_locks: StdMutex::new(std::collections::HashMap::new()),
|
|
driver_operation_locks: StdMutex::new(std::collections::HashMap::new()),
|
|
installation_operation_lock: tokio::sync::RwLock::new(()),
|
|
install_cancellations: Mutex::new(std::collections::HashMap::new()),
|
|
};
|
|
mgr.migrate_legacy_jre();
|
|
mgr.cleanup_pending_jre_dirs();
|
|
mgr.cleanup_orphan_jre_dirs();
|
|
mgr
|
|
}
|
|
|
|
/// Atomically load, modify, and persist the installation state.
|
|
///
|
|
/// Keep the closure free of async work: this lock exists specifically to
|
|
/// prevent one installation operation from saving an older snapshot over
|
|
/// another operation's successful update.
|
|
pub fn mutate_state<T>(&self, mutate: impl FnOnce(&mut AgentState) -> T) -> Result<T, String> {
|
|
let _guard = self.state_lock.lock().map_err(|_| "Agent installation state lock was poisoned".to_string())?;
|
|
let mut state = self.load_state();
|
|
let result = mutate(&mut state);
|
|
self.save_state(&state)?;
|
|
Ok(result)
|
|
}
|
|
|
|
/// Register a fresh cancellation token for an operation keyed by `key` (an
|
|
/// operation-scoped key; see the `agent_service` key helpers). Any existing
|
|
/// token for the same key is replaced; each install registers under a unique
|
|
/// operation id, so a stale cancelled token cannot leak into another
|
|
/// operation, and concurrent same-driver installs stay isolated.
|
|
pub async fn begin_install_cancellation(&self, key: &str) -> Arc<AgentInstallCancellation> {
|
|
let token = Arc::new(AgentInstallCancellation::new());
|
|
self.install_cancellations.lock().await.insert(key.to_string(), Arc::clone(&token));
|
|
token
|
|
}
|
|
|
|
/// Remove the token for `key` only if it is still `token` (guarded by
|
|
/// `Arc::ptr_eq`) so a concurrently registered fresh token is not removed.
|
|
pub async fn finish_install_cancellation(&self, key: &str, token: &Arc<AgentInstallCancellation>) {
|
|
let mut cancellations = self.install_cancellations.lock().await;
|
|
if cancellations.get(key).is_some_and(|current| Arc::ptr_eq(current, token)) {
|
|
cancellations.remove(key);
|
|
}
|
|
}
|
|
|
|
/// Request cancellation for the operation registered under `key`.
|
|
pub async fn cancel_install(&self, key: &str) {
|
|
if let Some(token) = self.install_cancellations.lock().await.get(key).cloned() {
|
|
token.cancel();
|
|
}
|
|
}
|
|
|
|
/// Whether the operation registered under `key` (an operation-scoped key,
|
|
/// see `agent_service` key helpers) has been cancelled.
|
|
pub async fn is_install_cancelled(&self, key: &str) -> bool {
|
|
self.install_cancellations.lock().await.get(key).is_some_and(|token| token.is_cancelled())
|
|
}
|
|
|
|
fn migrate_legacy_jre(&self) {
|
|
let legacy = self.base_dir.join("jre");
|
|
let versioned = self.jre_dir(DEFAULT_JRE_KEY);
|
|
if legacy.exists() && !versioned.exists() {
|
|
let _ = std::fs::rename(&legacy, &versioned);
|
|
}
|
|
}
|
|
|
|
/// Best-effort cleanup of `pending_jre_cleanup` paths recorded by previous
|
|
/// runs that fell back to renaming an old JRE aside on Windows. Successful
|
|
/// removals are pruned from the persisted state. Failures are kept for the
|
|
/// next launch and never block startup. (Issue #1100, D6.)
|
|
fn cleanup_pending_jre_dirs(&self) {
|
|
if self.load_state().pending_jre_cleanup.is_empty() {
|
|
return;
|
|
}
|
|
|
|
if let Err(err) = self.mutate_state(|state| {
|
|
let mut remaining = Vec::new();
|
|
for path in std::mem::take(&mut state.pending_jre_cleanup) {
|
|
if !path.exists() {
|
|
continue;
|
|
}
|
|
match std::fs::remove_dir_all(&path) {
|
|
Ok(()) => log::info!("Cleaned up pending JRE stash: {}", path.display()),
|
|
Err(err) => {
|
|
log::warn!("Pending JRE cleanup failed for {}: {err}", path.display());
|
|
remaining.push(path);
|
|
}
|
|
}
|
|
}
|
|
state.pending_jre_cleanup = remaining;
|
|
}) {
|
|
log::warn!("Failed to persist post-cleanup AgentState: {err}");
|
|
}
|
|
}
|
|
|
|
/// Sweep `base_dir` for orphan `*.old-*` JRE stash directories left behind
|
|
/// by previous runs (e.g. process crashed before the stash was recorded).
|
|
/// Best-effort — failures are ignored.
|
|
fn cleanup_orphan_jre_dirs(&self) {
|
|
let entries = match std::fs::read_dir(&self.base_dir) {
|
|
Ok(e) => e,
|
|
Err(_) => return,
|
|
};
|
|
for entry in entries.flatten() {
|
|
let Some(name) = entry.file_name().to_str().map(str::to_string) else {
|
|
continue;
|
|
};
|
|
// Match `<...>.old-<digits>-<...>` (typically `jre-21.old-...`),
|
|
// which is the suffix scheme used by stash_old_jre_dir.
|
|
if !name.starts_with("jre-") || !name.contains(".old-") {
|
|
continue;
|
|
}
|
|
let path = entry.path();
|
|
if !path.is_dir() {
|
|
continue;
|
|
}
|
|
match std::fs::remove_dir_all(&path) {
|
|
Ok(()) => log::info!("Cleaned up orphan JRE stash: {}", path.display()),
|
|
Err(err) => log::warn!("Orphan JRE cleanup failed for {}: {err}", path.display()),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn base_dir(&self) -> &PathBuf {
|
|
&self.base_dir
|
|
}
|
|
|
|
pub fn agent_app_version(&self) -> &str {
|
|
&self.app_version
|
|
}
|
|
|
|
pub fn jre_dir(&self, jre_key: &str) -> PathBuf {
|
|
self.base_dir.join(format!("jre-{jre_key}"))
|
|
}
|
|
|
|
pub fn jre_java_path(&self, jre_key: &str) -> PathBuf {
|
|
let dir = self.jre_dir(jre_key);
|
|
let java_name = if cfg!(windows) { "java.exe" } else { "java" };
|
|
let flat = dir.join("bin").join(java_name);
|
|
if flat.exists() {
|
|
return flat;
|
|
}
|
|
// Some macOS runtimes are unpacked with a Contents/Home/ layout.
|
|
let macos = dir.join("Contents").join("Home").join("bin").join(java_name);
|
|
if macos.exists() {
|
|
return macos;
|
|
}
|
|
flat
|
|
}
|
|
|
|
pub fn driver_dir(&self, db_type: &str) -> PathBuf {
|
|
self.base_dir.join("drivers").join(db_type)
|
|
}
|
|
|
|
pub fn driver_jar_path(&self, db_type: &str) -> PathBuf {
|
|
self.driver_dir(db_type).join("agent.jar")
|
|
}
|
|
|
|
pub fn driver_native_path(&self, db_type: &str) -> PathBuf {
|
|
let executable_name = if cfg!(windows) { "agent.exe" } else { "agent" };
|
|
self.driver_dir(db_type).join(executable_name)
|
|
}
|
|
|
|
pub fn driver_native_platform_path(&self, db_type: &str, platform: &str) -> PathBuf {
|
|
self.driver_dir(db_type).join(platform)
|
|
}
|
|
|
|
pub fn is_sqlite_worker_driver(db_type: &str) -> bool {
|
|
db_type == SQLITE_WORKER_DRIVER_KEY
|
|
}
|
|
|
|
pub fn driver_native_installed(&self, db_type: &str) -> bool {
|
|
if Self::is_sqlite_worker_driver(db_type) {
|
|
SQLITE_WORKER_NATIVE_PLATFORMS
|
|
.iter()
|
|
.all(|platform| self.driver_native_platform_path(db_type, platform).is_file())
|
|
} else {
|
|
self.driver_native_path(db_type).exists()
|
|
}
|
|
}
|
|
|
|
pub fn driver_launch_config_path(&self, db_type: &str) -> PathBuf {
|
|
self.driver_dir(db_type).join("agent-launch.json")
|
|
}
|
|
|
|
pub fn download_cache_dir(&self) -> PathBuf {
|
|
self.base_dir.join(DOWNLOAD_CACHE_DIR_NAME)
|
|
}
|
|
|
|
pub fn download_cache_max_age_days(&self) -> u64 {
|
|
DOWNLOAD_CACHE_MAX_AGE_DAYS
|
|
}
|
|
|
|
fn state_path(&self) -> PathBuf {
|
|
self.base_dir.join("state.json")
|
|
}
|
|
|
|
pub fn load_state(&self) -> AgentState {
|
|
std::fs::read_to_string(self.state_path())
|
|
.ok()
|
|
.and_then(|s| serde_json::from_str(strip_utf8_bom(&s)).ok())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
pub fn save_state(&self, state: &AgentState) -> Result<(), String> {
|
|
let dir = self.base_dir.clone();
|
|
std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
|
|
let json = serde_json::to_string_pretty(state).map_err(|e| e.to_string())?;
|
|
std::fs::write(self.state_path(), json).map_err(|e| e.to_string())
|
|
}
|
|
|
|
pub fn is_jre_installed(&self, jre_key: &str) -> bool {
|
|
std::fs::metadata(self.jre_java_path(jre_key)).is_ok_and(|metadata| metadata.is_file() && metadata.len() > 0)
|
|
}
|
|
|
|
pub fn is_driver_installed(&self, db_type: &str) -> bool {
|
|
self.is_driver_jar_valid(db_type)
|
|
|| self.driver_native_installed(db_type)
|
|
|| self.driver_launch_config_path(db_type).exists()
|
|
}
|
|
|
|
pub fn is_driver_jar_valid(&self, db_type: &str) -> bool {
|
|
is_valid_jar_file(&self.driver_jar_path(db_type))
|
|
}
|
|
|
|
pub fn driver_requires_java_runtime(&self, db_type: &str) -> bool {
|
|
self.is_driver_jar_valid(db_type)
|
|
&& !self.driver_native_installed(db_type)
|
|
&& !self.driver_launch_config_path(db_type).exists()
|
|
}
|
|
|
|
/// The JRE key an installed driver actually depends on, or `None` when the
|
|
/// installed artifact runs without a managed JRE (native binary or custom
|
|
/// launch config).
|
|
///
|
|
/// The on-disk artifact is the source of truth. The persisted `jre` string
|
|
/// is only consulted for genuine Java (jar) drivers, because it is written
|
|
/// unconditionally at install time and goes stale when a driver is
|
|
/// refactored from Java to a native language. Every "is this JRE still in
|
|
/// use?" question must route through here rather than reading `driver.jre`
|
|
/// directly, so stale metadata can never manufacture a false dependency.
|
|
pub fn installed_driver_jre_dependency<'a>(&self, state: &'a AgentState, db_type: &str) -> Option<&'a str> {
|
|
if !self.driver_requires_java_runtime(db_type) {
|
|
return None;
|
|
}
|
|
state.installed_drivers.get(db_type).map(|driver| driver.jre.as_str())
|
|
}
|
|
|
|
pub fn resolve_agent_launch_spec(
|
|
&self,
|
|
state: &AgentState,
|
|
driver_key: &str,
|
|
jre_key: &str,
|
|
) -> Result<AgentLaunchSpec, String> {
|
|
self.resolve_agent_launch_spec_with_extra_args(state, driver_key, jre_key, &[])
|
|
}
|
|
|
|
pub fn resolve_agent_launch_spec_with_extra_args(
|
|
&self,
|
|
state: &AgentState,
|
|
driver_key: &str,
|
|
jre_key: &str,
|
|
extra_java_args: &[String],
|
|
) -> Result<AgentLaunchSpec, String> {
|
|
if driver_key == "dameng" {
|
|
validate_dameng_java_system_properties(extra_java_args)?;
|
|
}
|
|
let driver_dir = self.driver_dir(driver_key);
|
|
let config_path = self.driver_launch_config_path(driver_key);
|
|
if config_path.exists() {
|
|
return self.resolve_configured_agent_launch_spec(driver_key, &driver_dir, &config_path);
|
|
}
|
|
|
|
let native_path = self.driver_native_path(driver_key);
|
|
if native_path.exists() {
|
|
let (native_path, driver_dir) = if native_path.is_relative() || driver_dir.is_relative() {
|
|
let native_path = native_path
|
|
.canonicalize()
|
|
.map_err(|e| format!("Failed to resolve {driver_key} native agent executable path: {e}"))?;
|
|
let driver_dir = driver_dir
|
|
.canonicalize()
|
|
.map_err(|e| format!("Failed to resolve {driver_key} native agent working directory: {e}"))?;
|
|
(native_path, driver_dir)
|
|
} else {
|
|
(native_path, driver_dir)
|
|
};
|
|
return Ok(AgentLaunchSpec::new(native_path).with_working_dir(driver_dir));
|
|
}
|
|
|
|
let jar_path = self.driver_jar_path(driver_key);
|
|
if jar_path.exists() {
|
|
let java = self.resolve_java_runtime(state, jre_key)?;
|
|
if !is_valid_jar_file(&jar_path) {
|
|
return Err(format!(
|
|
"{driver_key} driver jar is invalid or corrupt. Please reinstall it from the Driver Manager."
|
|
));
|
|
}
|
|
return Ok(AgentLaunchSpec::java_jar_with_extra_args(java, jar_path, extra_java_args));
|
|
}
|
|
|
|
Err(format!("{driver_key} driver is not installed. Please install it from the Driver Manager."))
|
|
}
|
|
|
|
fn resolve_configured_agent_launch_spec(
|
|
&self,
|
|
driver_key: &str,
|
|
driver_dir: &Path,
|
|
config_path: &Path,
|
|
) -> Result<AgentLaunchSpec, String> {
|
|
let json = std::fs::read_to_string(config_path)
|
|
.map_err(|e| format!("Failed to read {driver_key} agent launch config: {e}"))?;
|
|
let config: AgentLaunchConfig = serde_json::from_str(strip_utf8_bom(&json))
|
|
.map_err(|e| format!("Failed to parse {driver_key} agent launch config: {e}"))?;
|
|
let command = config.command.trim();
|
|
if command.is_empty() {
|
|
return Err(format!("{driver_key} agent launch config command is empty"));
|
|
}
|
|
let working_dir = config
|
|
.working_dir
|
|
.as_deref()
|
|
.map(|value| self.resolve_driver_launch_path(driver_dir, value))
|
|
.transpose()?
|
|
.unwrap_or_else(|| driver_dir.to_path_buf());
|
|
let program = self.resolve_driver_launch_path(driver_dir, command)?;
|
|
let args = config.args.iter().map(|arg| self.expand_agent_launch_template(driver_dir, arg)).collect::<Vec<_>>();
|
|
Ok(AgentLaunchSpec::new(program).with_args(args).with_working_dir(working_dir))
|
|
}
|
|
|
|
fn resolve_driver_launch_path(&self, driver_dir: &Path, value: &str) -> Result<PathBuf, String> {
|
|
let expanded = self.expand_agent_launch_template(driver_dir, value);
|
|
let path = PathBuf::from(&expanded);
|
|
if path.is_absolute() || expanded.contains('/') || expanded.contains('\\') || expanded.starts_with('.') {
|
|
return Ok(if path.is_absolute() { path } else { driver_dir.join(path) });
|
|
}
|
|
Ok(path)
|
|
}
|
|
|
|
fn expand_agent_launch_template(&self, driver_dir: &Path, value: &str) -> String {
|
|
value
|
|
.replace("{driver_dir}", &driver_dir.to_string_lossy())
|
|
.replace("{agent_dir}", &self.base_dir.to_string_lossy())
|
|
.replace("{platform}", Self::current_platform())
|
|
}
|
|
|
|
pub fn collect_driver_store_usage(&self, plugin_root: &Path) -> DriverStoreUsage {
|
|
let mut jres = Vec::new();
|
|
let mut jre_bytes = 0u64;
|
|
if let Ok(entries) = std::fs::read_dir(&self.base_dir) {
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
if !path.is_dir() {
|
|
continue;
|
|
}
|
|
let Some(name) = path.file_name().and_then(|v| v.to_str()) else {
|
|
continue;
|
|
};
|
|
if !name.starts_with("jre-") {
|
|
continue;
|
|
}
|
|
let key = name.trim_start_matches("jre-").to_string();
|
|
let bytes = path_size_bytes(&path);
|
|
jre_bytes = jre_bytes.saturating_add(bytes);
|
|
jres.push(DriverStoreUsageItem { id: key, bytes });
|
|
}
|
|
}
|
|
jres.sort_by(|left, right| left.id.cmp(&right.id));
|
|
|
|
let mut agent_drivers = Vec::new();
|
|
let mut agent_driver_bytes = 0u64;
|
|
let drivers_root = self.base_dir.join("drivers");
|
|
if let Ok(entries) = std::fs::read_dir(&drivers_root) {
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
if !path.is_dir() {
|
|
continue;
|
|
}
|
|
let Some(id) = path.file_name().and_then(|v| v.to_str()) else {
|
|
continue;
|
|
};
|
|
let bytes = path_size_bytes(&path);
|
|
agent_driver_bytes = agent_driver_bytes.saturating_add(bytes);
|
|
agent_drivers.push(DriverStoreUsageItem { id: id.to_string(), bytes });
|
|
}
|
|
}
|
|
agent_drivers.sort_by(|left, right| left.id.cmp(&right.id));
|
|
|
|
let jdbc_root = plugin_root.join("jdbc");
|
|
let jdbc_driver_root = jdbc_root.join("drivers");
|
|
let jdbc_driver_bytes = path_size_bytes(&jdbc_driver_root);
|
|
let jdbc_total_bytes = path_size_bytes(&jdbc_root);
|
|
let jdbc_plugin_bytes = jdbc_total_bytes.saturating_sub(jdbc_driver_bytes);
|
|
let download_cache_bytes = path_size_bytes(&self.download_cache_dir());
|
|
|
|
DriverStoreUsage {
|
|
total_bytes: jre_bytes
|
|
.saturating_add(agent_driver_bytes)
|
|
.saturating_add(download_cache_bytes)
|
|
.saturating_add(jdbc_plugin_bytes)
|
|
.saturating_add(jdbc_driver_bytes),
|
|
jre_bytes,
|
|
agent_driver_bytes,
|
|
download_cache_bytes,
|
|
jdbc_plugin_bytes,
|
|
jdbc_driver_bytes,
|
|
jres,
|
|
agent_drivers,
|
|
}
|
|
}
|
|
|
|
pub fn resolve_java_runtime(&self, state: &AgentState, jre_key: &str) -> Result<PathBuf, String> {
|
|
match state.java_runtime.mode {
|
|
JavaRuntimeMode::Managed => {
|
|
if !self.is_jre_installed(jre_key) {
|
|
return Err(format!(
|
|
"JRE {jre_key} runtime is not installed. Please install it from the Driver Manager."
|
|
));
|
|
}
|
|
Ok(self.jre_java_path(jre_key))
|
|
}
|
|
JavaRuntimeMode::System => {
|
|
let java_home = std::env::var_os(JAVA_HOME_ENV);
|
|
let path_var = std::env::var_os(PATH_ENV);
|
|
resolve_system_java_path(java_home.as_deref(), path_var.as_deref())
|
|
.ok_or_else(|| system_java_missing_error(java_home.as_deref(), path_var.as_deref()))
|
|
}
|
|
JavaRuntimeMode::Custom => {
|
|
let path = state
|
|
.java_runtime
|
|
.custom_java_path
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|path| !path.is_empty())
|
|
.ok_or_else(|| "Custom Java runtime path is empty. Please choose a Java executable.".to_string())?;
|
|
resolve_custom_java_path(path)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn stop_daemons(&self) {
|
|
crate::agent_runtime::stop_daemons(self).await;
|
|
}
|
|
|
|
pub async fn stop_daemon_by_key(&self, agent_key: &str) {
|
|
crate::agent_runtime::stop_daemon_by_key(self, agent_key).await;
|
|
}
|
|
|
|
pub async fn restart_daemon_by_key(&self, agent_key: &str) -> Result<(), String> {
|
|
crate::agent_runtime::restart_daemon_by_key(self, agent_key).await
|
|
}
|
|
|
|
pub async fn active_daemon_keys(&self) -> Vec<String> {
|
|
let mut keys = self.daemons.lock().await.keys().cloned().collect::<std::collections::HashSet<_>>();
|
|
for runtime_key in self.connection_runtimes.lock().await.keys() {
|
|
if let Some((agent_key, _)) = runtime_key.split_once('|') {
|
|
keys.insert(agent_key.to_string());
|
|
}
|
|
}
|
|
keys.into_iter().collect()
|
|
}
|
|
|
|
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
|
|
crate::agent_runtime::db_type_to_agent_key(db_type, driver_profile)
|
|
}
|
|
|
|
pub fn is_agent_type(db_type: &DatabaseType) -> bool {
|
|
crate::agent_runtime::is_agent_type(db_type)
|
|
}
|
|
|
|
pub async fn spawn(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
) -> Result<AgentDriverClient, String> {
|
|
self.spawn_with_extra_java_args(db_type, driver_profile, &[]).await
|
|
}
|
|
|
|
pub async fn spawn_with_extra_java_args(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
extra_java_args: &[String],
|
|
) -> Result<AgentDriverClient, String> {
|
|
crate::agent_runtime::spawn_connection_client(self, db_type, driver_profile, extra_java_args).await
|
|
}
|
|
|
|
pub async fn spawn_shared_connection_client(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
extra_java_args: &[String],
|
|
agent_session_id: String,
|
|
connect_params: serde_json::Value,
|
|
connect_timeout: std::time::Duration,
|
|
) -> Result<AgentDriverClient, crate::agent_runtime::SharedConnectionOpenError> {
|
|
crate::agent_runtime::spawn_shared_connection_client(
|
|
self,
|
|
db_type,
|
|
driver_profile,
|
|
extra_java_args,
|
|
agent_session_id,
|
|
connect_params,
|
|
connect_timeout,
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub async fn call_daemon<T: serde::de::DeserializeOwned + Send + 'static>(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
method: &str,
|
|
params: serde_json::Value,
|
|
) -> Result<T, String> {
|
|
crate::agent_runtime::call_daemon(self, db_type, driver_profile, method, params).await
|
|
}
|
|
|
|
pub async fn call_daemon_with_timeout<T: serde::de::DeserializeOwned + Send + 'static>(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
method: &str,
|
|
params: serde_json::Value,
|
|
timeout_duration: Option<std::time::Duration>,
|
|
) -> Result<T, String> {
|
|
crate::agent_runtime::call_daemon_with_timeout(self, db_type, driver_profile, method, params, timeout_duration)
|
|
.await
|
|
}
|
|
|
|
pub async fn call_daemon_method<T: serde::de::DeserializeOwned + Send + 'static>(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
method: AgentMethod,
|
|
params: serde_json::Value,
|
|
) -> Result<T, String> {
|
|
crate::agent_runtime::call_daemon_method(self, db_type, driver_profile, method, params).await
|
|
}
|
|
|
|
pub async fn call_daemon_method_with_timeout<T: serde::de::DeserializeOwned + Send + 'static>(
|
|
&self,
|
|
db_type: &DatabaseType,
|
|
driver_profile: Option<&str>,
|
|
method: AgentMethod,
|
|
params: serde_json::Value,
|
|
timeout_duration: Option<std::time::Duration>,
|
|
) -> Result<T, String> {
|
|
crate::agent_runtime::call_daemon_method_with_timeout(
|
|
self,
|
|
db_type,
|
|
driver_profile,
|
|
method,
|
|
params,
|
|
timeout_duration,
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub async fn download_file(url: &str, dest: &Path) -> Result<(), String> {
|
|
if let Some(parent) = dest.parent() {
|
|
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
|
|
}
|
|
let resp = reqwest::get(url).await.map_err(|e| format!("Download failed: {e}"))?;
|
|
if !resp.status().is_success() {
|
|
return Err(format!("Download failed with status: {}", resp.status()));
|
|
}
|
|
let bytes = resp.bytes().await.map_err(|e| format!("Download read failed: {e}"))?;
|
|
std::fs::write(dest, &bytes).map_err(|e| format!("Failed to write file: {e}"))
|
|
}
|
|
|
|
pub fn current_platform() -> &'static str {
|
|
if cfg!(target_os = "macos") && cfg!(target_arch = "aarch64") {
|
|
"macos-aarch64"
|
|
} else if cfg!(target_os = "macos") && cfg!(target_arch = "x86_64") {
|
|
"macos-x64"
|
|
} else if cfg!(target_os = "linux") && cfg!(target_arch = "aarch64") {
|
|
"linux-aarch64"
|
|
} else if cfg!(target_os = "linux") && cfg!(target_arch = "x86_64") {
|
|
"linux-x64"
|
|
} else if cfg!(target_os = "windows") && cfg!(target_arch = "aarch64") {
|
|
"windows-aarch64"
|
|
} else if cfg!(target_os = "windows") && cfg!(target_arch = "x86_64") {
|
|
"windows-x64"
|
|
} else {
|
|
"unknown"
|
|
}
|
|
}
|
|
}
|
|
|
|
fn path_size_bytes(path: &Path) -> u64 {
|
|
if let Ok(meta) = std::fs::symlink_metadata(path) {
|
|
if meta.is_file() {
|
|
return meta.len();
|
|
}
|
|
if !meta.is_dir() {
|
|
return 0;
|
|
}
|
|
} else {
|
|
return 0;
|
|
}
|
|
|
|
let mut total = 0u64;
|
|
if let Ok(entries) = std::fs::read_dir(path) {
|
|
for entry in entries.flatten() {
|
|
total = total.saturating_add(path_size_bytes(&entry.path()));
|
|
}
|
|
}
|
|
total
|
|
}
|
|
|
|
fn java_executable_name() -> &'static str {
|
|
if cfg!(windows) {
|
|
"java.exe"
|
|
} else {
|
|
"java"
|
|
}
|
|
}
|
|
|
|
/// Locate `java` under a JDK/JRE root. Accepts the conventional `bin/java`
|
|
/// layout, a macOS bundle (`Contents/Home/bin/java`), and a path that already
|
|
/// points at the executable itself.
|
|
fn resolve_java_under_root(root: &Path) -> Option<PathBuf> {
|
|
if is_executable_file(root) {
|
|
return Some(root.to_path_buf());
|
|
}
|
|
|
|
let flat = root.join("bin").join(java_executable_name());
|
|
if is_executable_file(&flat) {
|
|
return Some(flat);
|
|
}
|
|
|
|
let macos = root.join("Contents").join("Home").join("bin").join(java_executable_name());
|
|
if is_executable_file(&macos) {
|
|
return Some(macos);
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn resolve_custom_java_path(path: &str) -> Result<PathBuf, String> {
|
|
let raw = PathBuf::from(path);
|
|
resolve_java_under_root(&raw)
|
|
.ok_or_else(|| format!("Custom Java runtime does not exist or is not a Java executable: {}", raw.display()))
|
|
}
|
|
|
|
const JAVA_HOME_ENV: &str = "JAVA_HOME";
|
|
const PATH_ENV: &str = "PATH";
|
|
|
|
/// Resolve the "system" Java runtime from the user's own environment:
|
|
/// `JAVA_HOME` first (their explicit declaration), then the first `java` on
|
|
/// `PATH`. A `JAVA_HOME` that no longer contains a usable `java` is skipped
|
|
/// instead of reported, so a stale value cannot shadow a working `PATH` entry
|
|
/// (see the dangling-`JAVA_HOME` report in #5517).
|
|
///
|
|
/// The caller passes the environment explicitly so the tests stay hermetic.
|
|
fn resolve_system_java_path(java_home: Option<&OsStr>, path_var: Option<&OsStr>) -> Option<PathBuf> {
|
|
if let Some(home) = java_home.filter(|value| !value.is_empty()) {
|
|
if let Some(java) = resolve_java_under_root(Path::new(home)) {
|
|
return Some(java);
|
|
}
|
|
}
|
|
|
|
let path_var = path_var.filter(|value| !value.is_empty())?;
|
|
std::env::split_paths(path_var)
|
|
.filter(|dir| !dir.as_os_str().is_empty())
|
|
.map(|dir| dir.join(java_executable_name()))
|
|
.find(|candidate| is_executable_file(candidate))
|
|
}
|
|
|
|
/// Report where the "system" Java runtime was looked for. Users who configured
|
|
/// `JAVA_HOME`/`PATH` can still be running with the environment dbx was
|
|
/// launched with, so name the exact values that were checked instead of only
|
|
/// saying "not found".
|
|
fn system_java_missing_error(java_home: Option<&OsStr>, path_var: Option<&OsStr>) -> String {
|
|
let java_home_detail = match java_home.filter(|value| !value.is_empty()) {
|
|
Some(value) => format!(
|
|
"JAVA_HOME={} does not contain {}",
|
|
value.to_string_lossy(),
|
|
Path::new(value).join("bin").join(java_executable_name()).display()
|
|
),
|
|
None => "JAVA_HOME is not set".to_string(),
|
|
};
|
|
let path_detail = match path_var.filter(|value| !value.is_empty()) {
|
|
Some(value) => format!("no java in the {} PATH entries checked", std::env::split_paths(value).count()),
|
|
None => "PATH is not set".to_string(),
|
|
};
|
|
format!(
|
|
"System Java runtime was not found: {java_home_detail}; {path_detail}. \
|
|
Point JAVA_HOME at a JDK/JRE directory, add java to PATH, or choose a custom Java executable. \
|
|
dbx uses the environment it was started with, so restart dbx after changing environment variables."
|
|
)
|
|
}
|
|
|
|
fn is_executable_file(path: &Path) -> bool {
|
|
if !path.is_file() {
|
|
return false;
|
|
}
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
|
|
path.metadata().map(|meta| meta.permissions().mode() & 0o111 != 0).unwrap_or(false)
|
|
}
|
|
#[cfg(not(unix))]
|
|
{
|
|
true
|
|
}
|
|
}
|