diff --git a/crates/dbx-core/src/database_capabilities.rs b/crates/dbx-core/src/database_capabilities.rs new file mode 100644 index 000000000..0ba125d95 --- /dev/null +++ b/crates/dbx-core/src/database_capabilities.rs @@ -0,0 +1,53 @@ +use crate::models::connection::DatabaseType; + +pub fn agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> { + match db_type { + DatabaseType::Dameng => Some("dameng"), + DatabaseType::Kingbase => Some("kingbase"), + DatabaseType::Vastbase => Some("vastbase"), + DatabaseType::Goldendb => Some("goldendb"), + DatabaseType::Oracle => match driver_profile { + Some("oracle-10g") => Some("oracle-10g"), + _ => Some("oracle"), + }, + DatabaseType::H2 => Some("h2"), + DatabaseType::Snowflake => Some("snowflake"), + DatabaseType::Trino => Some("trino"), + DatabaseType::Hive => Some("hive"), + DatabaseType::Db2 => Some("db2"), + DatabaseType::Informix => Some("informix"), + DatabaseType::Neo4j => Some("neo4j"), + DatabaseType::Cassandra => Some("cassandra"), + DatabaseType::Bigquery => Some("bigquery"), + DatabaseType::Kylin => Some("kylin"), + DatabaseType::Sundb => Some("sundb"), + DatabaseType::Gaussdb => Some("gaussdb"), + _ => None, + } +} + +pub fn is_agent_type(db_type: &DatabaseType) -> bool { + agent_key(db_type, None).is_some() +} + +pub fn is_single_connection_pool(db_type: &DatabaseType) -> bool { + matches!( + db_type, + DatabaseType::Sqlite + | DatabaseType::DuckDb + | DatabaseType::Oracle + | DatabaseType::Dameng + | DatabaseType::Kingbase + | DatabaseType::Vastbase + | DatabaseType::Goldendb + | DatabaseType::Jdbc + ) +} + +pub fn is_metadata_connection_scoped(db_type: &DatabaseType) -> bool { + matches!(db_type, DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks) +} + +pub fn skips_tcp_probe(db_type: &DatabaseType) -> bool { + matches!(db_type, DatabaseType::Sqlite | DatabaseType::DuckDb | DatabaseType::Jdbc) || is_agent_type(db_type) +} diff --git a/crates/dbx-core/tests/database_capabilities.rs b/crates/dbx-core/tests/database_capabilities.rs new file mode 100644 index 000000000..136b5b7cb --- /dev/null +++ b/crates/dbx-core/tests/database_capabilities.rs @@ -0,0 +1,54 @@ +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; + +#[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::Gaussdb, None), Some("gaussdb")); + assert_eq!(agent_key(&DatabaseType::Oracle, None), Some("oracle")); + assert_eq!(agent_key(&DatabaseType::Oracle, Some("oracle-10g")), Some("oracle-10g")); + assert_eq!(agent_key(&DatabaseType::Postgres, None), None); +} + +#[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::Mysql)); + assert!(!is_agent_type(&DatabaseType::Jdbc)); +} + +#[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::Oracle)); + assert!(is_single_connection_pool(&DatabaseType::Dameng)); + assert!(is_single_connection_pool(&DatabaseType::Jdbc)); + assert!(!is_single_connection_pool(&DatabaseType::Trino)); + assert!(!is_single_connection_pool(&DatabaseType::Postgres)); +} + +#[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::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::Trino)); + assert!(skips_tcp_probe(&DatabaseType::Oracle)); + assert!(!skips_tcp_probe(&DatabaseType::Postgres)); + assert!(!skips_tcp_probe(&DatabaseType::Mysql)); +}