fix(schema): route Agent database enumeration through metadata pools

Refs: #8331

crates/dbx-core/src/schema/agent_metadata_fixture.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+import json
+import pathlib
+import sys
+
+root = next(parent for parent in pathlib.Path(__file__).parents if parent.name == 'agents')
+sessions = {}
+print(json.dumps({'ready': True}), flush=True)
+for line in sys.stdin:
+    request = json.loads(line)
+    with (root / 'requests.jsonl').open('a') as requests:
+        requests.write(json.dumps(request) + '\n')
+    method = request['method']
+    params = request.get('params', {})
+    session = params.get('agentSessionId')
+    result = {}
+    error = None
+    if method == 'handshake':
+        result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
+    elif method in ('open_session', 'connect'):
+        sessions[session] = params
+    elif method == 'close_session':
+        sessions.pop(session, None)
+    elif method == 'list_databases':
+        failure = root / 'list-error'
+        if failure.exists():
+            category = failure.read_text()
+            error = {'code': -1, 'message': 'fixture ' + category, 'data': {
+                'category': category, 'retryable': False,
+                'sessionDisposition': {'transport': 'replace_runtime', 'sql': 'keep'}.get(category, 'quarantine'),
+                'stage': 'execute'}}
+        else:
+            result = [{'name': sessions[session].get('sessionRole', 'workload')}]
+    response = {'jsonrpc': '2.0', 'id': request['id']}
+    response.update({'error': error} if error else {'result': result})
+    print(json.dumps(response), flush=True)

crates/dbx-core/src/schema/agent_metadata_routing_tests.rs

@@ -0,0 +1,233 @@
+use super::{list_databases_core, AppState, ConnectionConfig, DatabaseType, PoolKind};
+use crate::db::agent_driver::{AgentDriverClient, PooledAgentClient};
+use serde_json::{json, Value};
+use std::os::unix::fs::PermissionsExt;
+use std::sync::Arc;
+use std::time::Duration;
+
+struct AgentFixture {
+    state: Arc<AppState>,
+    directory: tempfile::TempDir,
+    _listener: tokio::net::TcpListener,
+}
+
+impl AgentFixture {
+    async fn new(db_type: DatabaseType) -> Self {
+        let directory = tempfile::tempdir().unwrap();
+        let storage = crate::persistence::test_storage::open(&directory.path().join("storage.db")).await.unwrap();
+        let state = Arc::new(AppState::new_with_plugin_and_agent_dir_and_app_version(
+            storage,
+            directory.path().join("plugins"),
+            directory.path().join("agents"),
+            "test",
+        ));
+        let profile = (db_type == DatabaseType::SqlServer).then_some("sqlserver-legacy");
+        let driver_key = crate::database_capabilities::agent_key(&db_type, profile).unwrap();
+        let executable = state.agent_manager.driver_native_path(driver_key);
+        std::fs::create_dir_all(executable.parent().unwrap()).unwrap();
+        std::fs::write(&executable, include_str!("agent_metadata_fixture.py")).unwrap();
+        std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap();
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let config: ConnectionConfig = serde_json::from_value(json!({
+            "id": "conn", "name": "Agent metadata fixture", "db_type": db_type,
+            "host": "127.0.0.1", "port": listener.local_addr().unwrap().port(),
+            "username": "fixture", "password": "", "database": "configured", "driver_profile": profile,
+            "connect_timeout_secs": 2, "query_timeout_secs": 1,
+            "keepalive_interval_secs": 0, "idle_timeout_secs": 0
+        }))
+        .unwrap();
+        state.configs.write().await.insert("conn".into(), config);
+        Self { state, directory, _listener: listener }
+    }
+
+    fn control_path(&self, name: &str) -> std::path::PathBuf {
+        let config = self.state.agent_manager.base_dir();
+        config.join(name)
+    }
+
+    fn requests(&self, method: &str) -> Vec<Value> {
+        std::fs::read_to_string(self.control_path("requests.jsonl"))
+            .unwrap_or_default()
+            .lines()
+            .map(|line| serde_json::from_str::<Value>(line).unwrap())
+            .filter(|request| request["method"] == method)
+            .collect()
+    }
+
+    async fn pool(&self, key: &str) -> Arc<PooledAgentClient> {
+        match self.state.pool_handle(key).await.unwrap() {
+            PoolKind::Agent(client) => client,
+            _ => panic!("expected Agent pool"),
+        }
+    }
+
+    async fn shutdown(self) {
+        self.state.shutdown(Duration::from_secs(2)).await;
+        drop(self.directory);
+    }
+}
+
+#[tokio::test]
+async fn enumeration_creates_metadata_pool_when_only_workload_exists() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    let workload_key = fixture.state.get_or_create_pool_for_session("conn", None, None).await.unwrap();
+    let workload = fixture.pool(&workload_key).await;
+
+    let databases = list_databases_core(&fixture.state, "conn").await.unwrap();
+
+    assert_eq!(databases[0].name, "metadata");
+    assert!(Arc::ptr_eq(&workload, &fixture.pool(&workload_key).await));
+    let opens = fixture.requests("open_session");
+    assert_eq!(opens.len(), 2);
+    assert_eq!(opens[1]["params"]["database"], "configured");
+    assert_eq!(fixture.requests("list_databases")[0]["params"]["agentSessionId"], opens[1]["params"]["agentSessionId"]);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_uses_metadata_pool_when_both_roles_exist() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    fixture.state.get_or_create_pool_for_session("conn", None, None).await.unwrap();
+    let metadata_key = fixture.state.get_or_create_metadata_pool_for_session("conn", None, None).await.unwrap();
+    let metadata = fixture.pool(&metadata_key).await;
+
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "metadata");
+    assert!(Arc::ptr_eq(&metadata, &fixture.pool(&metadata_key).await));
+    assert_eq!(fixture.requests("open_session").len(), 2);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_timeout_removes_metadata_and_next_call_opens_fresh_session() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    let workload_key = fixture.state.get_or_create_pool_for_session("conn", None, None).await.unwrap();
+    let workload = fixture.pool(&workload_key).await;
+    let metadata_key = fixture.state.get_or_create_metadata_pool_for_session("conn", None, None).await.unwrap();
+    let metadata = fixture.pool(&metadata_key).await;
+    std::fs::write(fixture.control_path("list-error"), "timeout").unwrap();
+
+    let error = list_databases_core(&fixture.state, "conn").await.unwrap_err();
+
+    assert!(error.contains("fixture timeout"), "{error}");
+    assert_eq!(fixture.requests("list_databases").len(), 1);
+    assert!(fixture.state.pool_handle(&metadata_key).await.is_none());
+    assert!(Arc::ptr_eq(&workload, &fixture.pool(&workload_key).await));
+    std::fs::remove_file(fixture.control_path("list-error")).unwrap();
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "metadata");
+    assert!(!Arc::ptr_eq(&metadata, &fixture.pool(&metadata_key).await));
+    assert_eq!(fixture.requests("list_databases").len(), 2);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_transport_failure_replaces_runtime_but_preserves_unrelated_workload() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    let unrelated = Arc::new(PooledAgentClient::new(AgentDriverClient::test_stub()));
+    fixture
+        .state
+        .update_connection_pools(|pools| {
+            pools.insert("conn".into(), PoolKind::Agent(unrelated.clone()));
+        })
+        .await;
+    let metadata_key = fixture.state.get_or_create_metadata_pool_for_session("conn", None, None).await.unwrap();
+    let failed = fixture.pool(&metadata_key).await;
+    std::fs::write(fixture.control_path("list-error"), "transport").unwrap();
+
+    let error = list_databases_core(&fixture.state, "conn").await.unwrap_err();
+
+    assert!(error.contains("fixture transport"), "{error}");
+    assert!(fixture.state.pool_handle(&metadata_key).await.is_none());
+    assert!(Arc::ptr_eq(&unrelated, &fixture.pool("conn").await));
+    assert_eq!(fixture.requests("list_databases").len(), 1);
+    std::fs::remove_file(fixture.control_path("list-error")).unwrap();
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "metadata");
+    assert!(!Arc::ptr_eq(&failed, &fixture.pool(&metadata_key).await));
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_native_sqlite_keeps_existing_pool() {
+    let directory = tempfile::tempdir().unwrap();
+    let storage = crate::persistence::test_storage::open(&directory.path().join("storage.db")).await.unwrap();
+    let state = AppState::new(storage);
+    let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
+    state
+        .update_connection_pools(|pools| {
+            pools.insert("sqlite".into(), PoolKind::Sqlite(pool));
+        })
+        .await;
+    assert!(!list_databases_core(&state, "sqlite").await.unwrap().is_empty());
+    assert!(state.pool_handle("sqlite").await.is_some());
+    state.shutdown(Duration::from_secs(2)).await;
+}
+
+#[tokio::test]
+async fn enumeration_ignores_bare_client_with_unavailable_stdin() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    fixture
+        .state
+        .update_connection_pools(|pools| {
+            pools.insert("conn".into(), PoolKind::agent(AgentDriverClient::test_stub()));
+        })
+        .await;
+
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "metadata");
+    assert_eq!(fixture.requests("open_session").len(), 1);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_runtime_fail_stop_removes_shared_workload_and_recreates_without_bare_pool() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    let workload_key = fixture.state.get_or_create_pool_for_session("conn", None, None).await.unwrap();
+    let metadata_key = fixture.state.get_or_create_metadata_pool_for_session("conn", None, None).await.unwrap();
+    std::fs::write(fixture.control_path("list-error"), "transport").unwrap();
+
+    assert!(list_databases_core(&fixture.state, "conn").await.unwrap_err().contains("fixture transport"));
+    assert!(fixture.state.pool_handle(&metadata_key).await.is_none());
+    assert!(fixture.state.pool_handle(&workload_key).await.is_none());
+    std::fs::remove_file(fixture.control_path("list-error")).unwrap();
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "metadata");
+    assert!(fixture.state.pool_handle(&workload_key).await.is_none());
+    assert_eq!(fixture.requests("open_session").len(), 3);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_sql_error_keeps_metadata_session_without_replay() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    let metadata_key = fixture.state.get_or_create_metadata_pool_for_session("conn", None, None).await.unwrap();
+    let metadata = fixture.pool(&metadata_key).await;
+    std::fs::write(fixture.control_path("list-error"), "sql").unwrap();
+
+    assert!(list_databases_core(&fixture.state, "conn").await.unwrap_err().contains("fixture sql"));
+    assert!(Arc::ptr_eq(&metadata, &fixture.pool(&metadata_key).await));
+    assert_eq!(fixture.requests("list_databases").len(), 1);
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_mongo_keeps_document_workload_route() {
+    let fixture = AgentFixture::new(DatabaseType::Hive).await;
+    fixture.state.get_or_create_pool_for_session("conn", None, None).await.unwrap();
+    fixture.state.configs.write().await.get_mut("conn").unwrap().db_type = DatabaseType::MongoDb;
+
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "workload");
+    assert_eq!(fixture.requests("open_session").len(), 1);
+    assert!(fixture.state.pool_handle("conn:role:metadata").await.is_none());
+    fixture.shutdown().await;
+}
+
+#[tokio::test]
+async fn enumeration_legacy_sqlserver_reuses_connection_metadata_across_databases() {
+    let fixture = AgentFixture::new(DatabaseType::SqlServer).await;
+    let metadata_key =
+        fixture.state.get_or_create_metadata_pool_for_session("conn", Some("other"), None).await.unwrap();
+    let metadata = fixture.pool(&metadata_key).await;
+
+    assert_eq!(list_databases_core(&fixture.state, "conn").await.unwrap()[0].name, "workload");
+    assert!(Arc::ptr_eq(&metadata, &fixture.pool(&metadata_key).await));
+    assert_eq!(fixture.requests("connect").len(), 1);
+    assert_eq!(metadata_key, "conn:role:metadata");
+    fixture.shutdown().await;
+}

crates/dbx-core/src/schema/external_table_filter_tests.rs

@@ -19,7 +19,9 @@ tables = [{'name': 'T%04d' % i, 'table_type': 'TABLE'} for i in range(1005)]
 for line in sys.stdin:
     request = json.loads(line)
     params = request.get('params', {})
-    if request['method'] == 'listTables':
+    if request['method'] == 'listDatabases':
+        result = [{'name': params['connection']['database']}]
+    elif request['method'] == 'listTables':
         with pathlib.Path(__file__).with_name('calls.jsonl').open('a') as log:
             log.write(json.dumps(params) + '\n')
         start = params.get('offset', 0)
@@ -77,6 +79,18 @@ for line in sys.stdin:
     (dir, state)
 }
 
+#[tokio::test]
+async fn enumeration_preserves_external_jdbc_plugin_even_with_agent_vendor_config() {
+    let (_directory, state) = jdbc_state().await;
+    state.configs.write().await.get_mut("jdbc-tables").unwrap().db_type = crate::models::connection::DatabaseType::Hive;
+
+    let databases = super::list_databases_core(&state, "jdbc-tables").await.unwrap();
+
+    assert_eq!(databases[0].name, "demo");
+    assert!(state.pool_handle("jdbc-tables:role:metadata").await.is_none());
+    state.shutdown(Duration::from_secs(2)).await;
+}
+
 #[tokio::test]
 async fn jdbc_table_filters_find_matches_beyond_the_plugin_page_and_page_after_filtering() {
     let (dir, state) = jdbc_state().await;

crates/dbx-core/src/schema/mod.rs

@@ -17,6 +17,8 @@ use std::future::Future;
 use std::sync::{Arc, OnceLock};
 use std::time::{Duration, Instant};
 
+#[cfg(all(test, unix))]
+mod agent_metadata_routing_tests;
 mod agent_pg_sequences;
 #[cfg(all(test, unix))]
 mod external_table_filter_tests;
@@ -688,12 +690,21 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result<Ve
             return Ok(vec![db::DatabaseInfo { name, ..Default::default() }]);
         }
         try_sqlserver!(pool_handle, list_databases);
-        if let Some(client) = extract_pool!(pool_handle.as_ref(), Agent) {
+        if matches!(pool_handle.as_ref(), Some(PoolKind::Agent(_)))
+            || (pool_handle.is_none()
+                && db_config.as_ref().is_some_and(|config| {
+                    crate::database_capabilities::is_agent_type(&config.db_type)
+                        || crate::connection::sqlserver_uses_legacy_driver(config)
+                }))
+        {
             let is_mongo = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::MongoDb);
             if is_mongo {
                 let dbs = crate::mongo_ops::mongo_list_databases_core(state, connection_id).await?;
                 return Ok(dbs.into_iter().map(|name| db::DatabaseInfo { name, ..Default::default() }).collect());
             }
+            let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
+            let metadata_pool = state.pool_handle(&pool_key).await;
+            let client = extract_pool!(metadata_pool.as_ref(), Agent).ok_or("Pool not found")?;
             let mut client = client.lock().await;
             return client.list_databases(agent_metadata_timeout(db_config.as_ref())).await;
         }