diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 336c18e7e..a05671ff5 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -31,6 +31,47 @@ pub struct AgentHandshake { pub capabilities: Vec, } +impl AgentHandshake { + pub fn supports(&self, capability: AgentCapability) -> bool { + self.capabilities.iter().any(|value| value == capability.as_str()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentCapability { + Connect, + TestConnection, + Metadata, + Query, + PagedQuery, + Transaction, + Ddl, +} + +impl AgentCapability { + pub const ALL: [Self; 7] = [ + Self::Connect, + Self::TestConnection, + Self::Metadata, + Self::Query, + Self::PagedQuery, + Self::Transaction, + Self::Ddl, + ]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Connect => "connect", + Self::TestConnection => "test_connection", + Self::Metadata => "metadata", + Self::Query => "query", + Self::PagedQuery => "paged_query", + Self::Transaction => "transaction", + Self::Ddl => "ddl", + } + } +} + struct StderrTail { lines: VecDeque, capacity: usize, @@ -378,7 +419,8 @@ impl Drop for AgentDriverClient { mod tests { use super::{ agent_handshake_params, agent_java_args, agent_proxy_env_vars, format_agent_process_error, - is_unsupported_handshake_error, read_agent_line, AgentHandshake, StderrTail, AGENT_PROTOCOL_VERSION, + is_unsupported_handshake_error, read_agent_line, AgentCapability, AgentHandshake, StderrTail, + AGENT_PROTOCOL_VERSION, }; use std::io::Cursor; @@ -473,6 +515,31 @@ mod tests { assert_eq!(handshake.capabilities, vec!["connect", "query", "metadata"]); } + #[test] + fn defines_agent_protocol_capabilities() { + assert_eq!(AgentCapability::Connect.as_str(), "connect"); + assert_eq!(AgentCapability::TestConnection.as_str(), "test_connection"); + assert_eq!(AgentCapability::Metadata.as_str(), "metadata"); + assert_eq!(AgentCapability::Query.as_str(), "query"); + assert_eq!(AgentCapability::PagedQuery.as_str(), "paged_query"); + assert_eq!(AgentCapability::Transaction.as_str(), "transaction"); + assert_eq!(AgentCapability::Ddl.as_str(), "ddl"); + assert_eq!(AgentCapability::ALL.len(), 7); + } + + #[test] + fn checks_handshake_capability_support() { + let handshake = AgentHandshake { + protocol_version: AGENT_PROTOCOL_VERSION, + agent_protocol_version: AGENT_PROTOCOL_VERSION, + capabilities: vec!["connect".to_string(), "metadata".to_string()], + }; + + assert!(handshake.supports(AgentCapability::Connect)); + assert!(handshake.supports(AgentCapability::Metadata)); + assert!(!handshake.supports(AgentCapability::Query)); + } + #[test] fn treats_unknown_handshake_method_as_compatible_fallback() { assert!(is_unsupported_handshake_error("Agent RPC error (-1): Unknown method: handshake"));