use dbx_core::agent_catalog; use dbx_core::database_capabilities::{ agent_key, is_agent_type, is_metadata_connection_scoped, is_single_connection_pool, skips_tcp_probe, }; use dbx_core::models::connection::DatabaseType; use serde::Deserialize; use std::collections::HashSet; #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct DriverManifest { drivers: Vec, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct DriverManifestEntry { db_type: DatabaseType, label: String, runtime_mode: String, support_level: String, capabilities: DriverProductCapabilities, #[serde(default)] agent_key: Option, #[serde(default)] single_connection_pool: bool, #[serde(default)] metadata_connection_scoped: bool, #[serde(default)] skip_tcp_probe: bool, #[serde(default)] driver_profiles: Vec, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct DriverProductCapabilities { query_execution: bool, metadata_browse: bool, object_browser: bool, object_source: bool, schema_search: bool, diagram: bool, table_data_edit: bool, table_structure_edit: bool, table_import: bool, data_transfer: bool, sql_file_execution: bool, database_create: bool, field_lineage: bool, sql_explain: bool, user_admin: bool, driver_management: bool, } impl DriverProductCapabilities { fn any_enabled(&self) -> bool { [ self.metadata_browse, self.query_execution, self.object_browser, self.object_source, self.schema_search, self.diagram, self.table_data_edit, self.table_structure_edit, self.table_import, self.data_transfer, self.sql_file_execution, self.database_create, self.field_lineage, self.sql_explain, self.user_admin, self.driver_management, ] .into_iter() .any(|enabled| enabled) } } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct DriverProfileEntry { profile: String, agent_key: String, } fn driver_manifest() -> DriverManifest { let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("assets").join("database-drivers.manifest.json"); let json = std::fs::read_to_string(&path).unwrap_or_else(|err| { panic!("failed to read driver manifest {}: {err}", path.display()); }); serde_json::from_str(&json).expect("driver manifest should be valid JSON") } #[test] fn maps_agent_database_types_to_driver_keys() { assert_eq!(agent_key(&DatabaseType::Trino, None), Some("trino")); assert_eq!(agent_key(&DatabaseType::Hive, None), Some("hive")); assert_eq!(agent_key(&DatabaseType::Tdengine, None), Some("tdengine")); assert_eq!(agent_key(&DatabaseType::Iotdb, None), Some("iotdb")); assert_eq!(agent_key(&DatabaseType::Yashandb, None), Some("yashandb")); assert_eq!(agent_key(&DatabaseType::Databricks, None), Some("databricks")); assert_eq!(agent_key(&DatabaseType::SapHana, None), Some("saphana")); assert_eq!(agent_key(&DatabaseType::Teradata, None), Some("teradata")); assert_eq!(agent_key(&DatabaseType::Vertica, None), Some("vertica")); assert_eq!(agent_key(&DatabaseType::Firebird, None), Some("firebird")); assert_eq!(agent_key(&DatabaseType::Exasol, None), Some("exasol")); assert_eq!(agent_key(&DatabaseType::OceanbaseOracle, None), Some("oceanbase-oracle")); assert_eq!(agent_key(&DatabaseType::Gbase, None), Some("gbase8a")); assert_eq!(agent_key(&DatabaseType::Access, None), Some("access")); assert_eq!(agent_key(&DatabaseType::Oracle, None), Some("oracle")); assert_eq!(agent_key(&DatabaseType::Databend, None), Some("databend")); assert_eq!(agent_key(&DatabaseType::InfluxDb, None), Some("influxdb")); assert_eq!(agent_key(&DatabaseType::ZooKeeper, None), Some("zookeeper")); assert_eq!(agent_key(&DatabaseType::Oracle, Some("oracle-legacy")), Some("oracle")); assert_eq!(agent_key(&DatabaseType::Oracle, Some("oracle-10g")), Some("oracle")); assert_eq!(agent_key(&DatabaseType::SqlServer, Some("sqlserver-legacy")), Some("sqlserver-legacy")); assert_eq!(agent_key(&DatabaseType::SqlServer, None), None); assert_eq!(agent_key(&DatabaseType::Postgres, None), None); } #[test] fn driver_store_entries_do_not_repeat_agent_keys() { let entries: Vec<_> = agent_catalog::driver_store_entries().collect(); let mut seen = HashSet::new(); let duplicate_keys: Vec<_> = entries.iter().map(|(key, _)| *key).filter(|key| !seen.insert(*key)).collect(); assert!(duplicate_keys.is_empty(), "driver store agent keys should be unique: {duplicate_keys:?}"); assert_eq!(entries.iter().filter(|(key, _)| *key == "gbase8a").count(), 1); assert_eq!(entries.iter().filter(|(key, _)| *key == "gbase8s").count(), 1); assert_eq!(entries.iter().filter(|(key, _)| *key == "sqlserver-legacy").count(), 1); assert_eq!(agent_catalog::label_for_key("sqlserver-legacy"), Some("SQL Server legacy compatibility component")); } #[test] fn classifies_agent_database_types() { assert!(is_agent_type(&DatabaseType::Oracle)); assert!(is_agent_type(&DatabaseType::Trino)); assert!(is_agent_type(&DatabaseType::Hive)); assert!(is_agent_type(&DatabaseType::Tdengine)); assert!(is_agent_type(&DatabaseType::Iotdb)); assert!(is_agent_type(&DatabaseType::Yashandb)); assert!(is_agent_type(&DatabaseType::Databricks)); assert!(is_agent_type(&DatabaseType::SapHana)); assert!(is_agent_type(&DatabaseType::Teradata)); assert!(is_agent_type(&DatabaseType::Vertica)); assert!(is_agent_type(&DatabaseType::Firebird)); assert!(is_agent_type(&DatabaseType::Exasol)); assert!(is_agent_type(&DatabaseType::OceanbaseOracle)); assert!(is_agent_type(&DatabaseType::Gbase)); assert!(is_agent_type(&DatabaseType::Access)); assert!(is_agent_type(&DatabaseType::Databend)); assert!(is_agent_type(&DatabaseType::InfluxDb)); assert!(is_agent_type(&DatabaseType::ZooKeeper)); assert!(!is_agent_type(&DatabaseType::Mysql)); assert!(!is_agent_type(&DatabaseType::Jdbc)); assert!(!is_agent_type(&DatabaseType::Gaussdb)); assert!(!is_agent_type(&DatabaseType::Kwdb)); assert!(!is_agent_type(&DatabaseType::OpenGauss)); assert!(!is_agent_type(&DatabaseType::Questdb)); } #[test] fn identifies_single_connection_pool_types() { assert!(is_single_connection_pool(&DatabaseType::Sqlite)); assert!(is_single_connection_pool(&DatabaseType::DuckDb)); assert!(is_single_connection_pool(&DatabaseType::MongoDb)); assert!(is_single_connection_pool(&DatabaseType::Oracle)); assert!(is_single_connection_pool(&DatabaseType::Dameng)); assert!(is_single_connection_pool(&DatabaseType::Access)); assert!(is_single_connection_pool(&DatabaseType::Yashandb)); assert!(is_single_connection_pool(&DatabaseType::Firebird)); assert!(is_single_connection_pool(&DatabaseType::OceanbaseOracle)); assert!(is_single_connection_pool(&DatabaseType::Jdbc)); assert!(!is_single_connection_pool(&DatabaseType::Trino)); assert!(!is_single_connection_pool(&DatabaseType::Postgres)); assert!(!is_single_connection_pool(&DatabaseType::Kwdb)); } #[test] fn identifies_metadata_connections_that_drop_database_scope() { assert!(is_metadata_connection_scoped(&DatabaseType::Mysql)); assert!(!is_metadata_connection_scoped(&DatabaseType::Doris)); assert!(!is_metadata_connection_scoped(&DatabaseType::StarRocks)); assert!(!is_metadata_connection_scoped(&DatabaseType::Postgres)); assert!(!is_metadata_connection_scoped(&DatabaseType::Kwdb)); assert!(!is_metadata_connection_scoped(&DatabaseType::Oracle)); } #[test] fn skips_tcp_probe_for_local_file_plugin_and_agent_types() { assert!(skips_tcp_probe(&DatabaseType::Sqlite)); assert!(skips_tcp_probe(&DatabaseType::DuckDb)); assert!(skips_tcp_probe(&DatabaseType::Jdbc)); assert!(skips_tcp_probe(&DatabaseType::Access)); assert!(skips_tcp_probe(&DatabaseType::H2)); assert!(skips_tcp_probe(&DatabaseType::Trino)); assert!(skips_tcp_probe(&DatabaseType::Oracle)); assert!(skips_tcp_probe(&DatabaseType::Tdengine)); assert!(skips_tcp_probe(&DatabaseType::Iotdb)); assert!(skips_tcp_probe(&DatabaseType::Yashandb)); assert!(skips_tcp_probe(&DatabaseType::Databricks)); assert!(skips_tcp_probe(&DatabaseType::OceanbaseOracle)); assert!(skips_tcp_probe(&DatabaseType::Gbase)); assert!(skips_tcp_probe(&DatabaseType::Databend)); assert!(skips_tcp_probe(&DatabaseType::InfluxDb)); assert!(skips_tcp_probe(&DatabaseType::MessageQueue)); assert!(skips_tcp_probe(&DatabaseType::ZooKeeper)); assert!(!skips_tcp_probe(&DatabaseType::Postgres)); assert!(!skips_tcp_probe(&DatabaseType::Mysql)); assert!(!skips_tcp_probe(&DatabaseType::Gaussdb)); assert!(!skips_tcp_probe(&DatabaseType::Kwdb)); assert!(!skips_tcp_probe(&DatabaseType::OpenGauss)); assert!(!skips_tcp_probe(&DatabaseType::Questdb)); } #[test] fn driver_manifest_matches_core_database_capabilities() { let manifest = driver_manifest(); let support_levels = ["connect", "browse", "understand", "operate"]; for driver in &manifest.drivers { // MQ is a message queue, not a database — skip database capability checks if driver.db_type == DatabaseType::MessageQueue { continue; } assert!( support_levels.contains(&driver.support_level.as_str()), "invalid support level for {:?}", driver.db_type ); assert!( driver.capabilities.any_enabled(), "database {:?} should declare at least one product capability", driver.db_type ); assert_eq!( is_agent_type(&driver.db_type), driver.runtime_mode == "agent", "agent classification mismatch for {:?}", driver.db_type ); assert_eq!( agent_key(&driver.db_type, None), driver.agent_key.as_deref(), "agent key mismatch for {:?}", driver.db_type ); assert_eq!( is_single_connection_pool(&driver.db_type), driver.single_connection_pool, "single-pool mismatch for {:?}", driver.db_type ); assert_eq!( is_metadata_connection_scoped(&driver.db_type), driver.metadata_connection_scoped, "metadata scope mismatch for {:?}", driver.db_type ); assert_eq!( skips_tcp_probe(&driver.db_type), driver.skip_tcp_probe, "TCP probe behavior mismatch for {:?}", driver.db_type ); for profile in &driver.driver_profiles { assert_eq!( agent_key(&driver.db_type, Some(&profile.profile)), Some(profile.agent_key.as_str()), "profile agent key mismatch for {:?}/{}", driver.db_type, profile.profile ); } } } #[test] fn driver_manifest_declares_expected_product_capabilities() { let manifest = driver_manifest(); let find_driver = |db_type: DatabaseType| { manifest .drivers .iter() .find(|driver| driver.db_type == db_type) .unwrap_or_else(|| panic!("{db_type:?} manifest entry")) }; let mysql = find_driver(DatabaseType::Mysql); assert_eq!(mysql.support_level, "operate"); assert!(mysql.capabilities.schema_search); assert!(mysql.capabilities.table_structure_edit); assert!(mysql.capabilities.sql_explain); assert!(!mysql.capabilities.driver_management); let jdbc = find_driver(DatabaseType::Jdbc); assert_eq!(jdbc.support_level, "browse"); assert!(jdbc.capabilities.metadata_browse); assert!(jdbc.capabilities.object_browser); assert!(jdbc.capabilities.sql_file_execution); assert!(!jdbc.capabilities.table_structure_edit); assert!(!jdbc.capabilities.user_admin); let manticore = find_driver(DatabaseType::ManticoreSearch); assert_eq!(manticore.support_level, "operate"); assert!(manticore.capabilities.metadata_browse); assert!(manticore.capabilities.sql_file_execution); assert!(manticore.capabilities.table_structure_edit); assert!(!manticore.capabilities.object_browser); assert!(manticore.capabilities.table_data_edit); let redis = find_driver(DatabaseType::Redis); assert_eq!(redis.support_level, "connect"); assert!(!redis.capabilities.object_browser); assert!(!redis.capabilities.sql_file_execution); let zookeeper = find_driver(DatabaseType::ZooKeeper); assert_eq!(zookeeper.label, "Apache ZooKeeper"); assert_eq!(zookeeper.runtime_mode, "agent"); assert_eq!(zookeeper.agent_key.as_deref(), Some("zookeeper")); assert_eq!(zookeeper.support_level, "connect"); assert!(zookeeper.capabilities.query_execution); assert!(zookeeper.capabilities.driver_management); assert!(!zookeeper.capabilities.metadata_browse); } #[test] fn kingbase_declares_data_transfer_support() { let manifest = driver_manifest(); let kingbase = manifest .drivers .iter() .find(|driver| driver.db_type == DatabaseType::Kingbase) .expect("Kingbase manifest entry"); assert!(kingbase.capabilities.table_import); assert!(kingbase.capabilities.data_transfer); }