refactor(core): centralize database capabilities

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)
+}

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));
+}